package service

import (
	"fmt"
	"gorm.io/gorm"
	"shyxy-net/common/cache/lru"
	"shyxy-net/common/db/model-service/model"
	"shyxy-net/model-service/repository"
)

type ModelService struct {
	modelInfoRepo     repostiroy.ModelInfoRepository
	rpcServerInfoRepo repostiroy.RpcServerInfoRepository
	lru               *lru.LRU
}

func New(modelInfoDB *gorm.DB, rpcServerDB *gorm.DB, cache *lru.LRU) *ModelService {
	cachedModelInfoRepo := repostiroy.NewCachedModelInfoRepository(modelInfoDB, lru.New(1000))
	cachedRpcServerRepo := repostiroy.NewCachedRpcServerInfoRepository(rpcServerDB, lru.New(1000))
	return &ModelService{
		modelInfoRepo:     cachedModelInfoRepo,
		rpcServerInfoRepo: cachedRpcServerRepo,
		lru:               cache,
	}
}

func (ms *ModelService) ListModels() ([]string, error) {
	modelInfos, err := ms.modelInfoRepo.ListModelInfos()
	if err != nil {
		return nil, err
	}
	if len(modelInfos) == 0 {
		return nil, fmt.Errorf("no available model found")
	}
	modelNames := make([]string, len(modelInfos))
	for i, modelInfo := range modelInfos {
		modelNames[i] = modelInfo.Name
	}
	return modelNames, nil
}

func (ms *ModelService) GetRpcServerAddressByModelName(modelName string) (string, error) {
	var err error
	address, err := computeRpcAddr(modelName, ms.rpcServerInfoRepo, ms.lru)
	if err != nil {
		return "", err
	}
	return address, nil
}

func (ms *ModelService) CheckIfModelExists(modelName string) (bool, error) {
	_, err := ms.modelInfoRepo.FindModelInfoByModelName(modelName)
	if err != nil {
		return false, err
	}
	return true, nil
}

// TODO 分配规则还有些错乱
func computeRpcAddr(modelName string, r repostiroy.RpcServerInfoRepository, cache *lru.LRU) (string, error) {
	availableRpcServerInfos, err := r.FindRpcServerInfosByModelName(modelName)
	if err != nil {
		return "", err
	}
	if len(availableRpcServerInfos) == 0 {
		return "", fmt.Errorf("no available rpc server")
	}
	serverInfo := &model.RpcServerInfo{
		ModelName:   modelName,
		ServiceAddr: availableRpcServerInfos[0].ServiceAddr,
	}
	generateCacheKey := func(modelName, addr string) string {
		return fmt.Sprintf("rpc_addr_%s_%s", modelName, addr)
	}
	backendAddrs := make([]string, len(availableRpcServerInfos))
	for i, info := range availableRpcServerInfos {
		backendAddrs[i] = info.ServiceAddr
		cacheKey := generateCacheKey(modelName, info.ServiceAddr)
		cache.Cache.Add(cacheKey, info.ServiceAddr)
	}
	cache.Add(fmt.Sprintf("rpc_addrs_%s", modelName), backendAddrs)
	return serverInfo.ServiceAddr, nil
}
